import dashscope
import os
import json
import sys
from openai import OpenAI
from loguru import logger
from http import HTTPStatus
sys.path.append(os.path.abspath('../LLM_APIs'))
from response_process import eval_response, process_llm_history


def call_with_messages(prompt, q_model='qwen-turbo', temperature=0, top_p=0, max_tokens=200, api_key=None):
    dashscope.api_key = api_key
    try:
        messages = [{'role': 'system', 'content': 'You are a helpful assistant.'},
                    {'role': 'user', 'content': prompt}]

        response = dashscope.Generation.call(
            model = q_model,  # 模型选择 qwen-max qwen-turbo qwen-plus
            messages = messages,
            # max_tokens = max_tokens,   # 最大字数限制
            top_p = top_p,  # 多样性设置
            temperature = temperature,  # 随机性设置
            repetition_penalty = 1.2,   # 重复惩罚设置
            result_format = 'message'  # 结果格式
        )
        return response
    except Exception as e:
        print(f"call Qwen error: {e}")
        return None


def build_qwen_prompt(model='qwen-turbo', **kwargs):
    try:
        model = kwargs['paras']['model']
    except:
        pass

    try:
        api_key = kwargs['config']['QWEN_APIKEY']
        task = kwargs['paras']['task']
        query = kwargs['paras']['query']
        texts = kwargs['paras']['texts']
    except Exception as e:
        print('Bad paras inputs... because: {}'.format(e))
        return None
    
    his_record = process_llm_history(kwargs['config'], kwargs['paras'])
    print(his_record)
    temperature = 0.01
    top_p = 0.01
    max_tokens = 1500
    
    if task=='talk': # when generating response using KB
        add_req = kwargs['paras']['add_req']
        max_tokens = kwargs['paras']['out_limit']
        temperature = 0
        top_p = 0.01
        
        prompt = f"""
                你将接收到用户的问题：'''{query}'''
                你可能还会接收到历史对话中的相关问题和答案：'''{his_record}'''
                请严格参考历史记录（如有），回答用户问题。{add_req}
                你的返回只能是json形式，键为"answer"，值为你的回答，不要返回其他信息。
            """
            
    elif task=='talk-kb':
        add_req = kwargs['paras']['add_req']
        if_judge = kwargs['paras']['if_judge']
        max_tokens = kwargs['paras']['out_limit']
        temperature = 0
        top_p = 0.01
        
        if if_judge:
            example = '''
                    {
                        "judge": "否",
                        "answer": "基础资料暂无答案..."
                    }
                '''
            
            prompt = f"""
                    你将接收到用户的问题：'''{query}'''
                    你将接收到一些基础知识和资料：'''{texts}'''
                    你可能还会接收到历史对话中的相关问题和答案：'''{his_record}'''
                    你要执行以下两步：
                    1）判断基础知识资料和历史回答中是否涉及用户问题的答案。
                    2）根据基础知识资料以及历史记录回答用户问题。{add_req}
                    
                    你的返回必须是JSON字典格式，包含2个键，第1个键是"judge"，值为基础知识资料或历史回答是否包含了用户问题的答案，只能“是”或“否"；第2个键是"answer"，值为你的回答，下面是一个返回示例
                """    
            prompt = prompt + example
        else:
            prompt = f"""
                    你将接收到用户的问题：'''{query}'''
                    你将接收到一些基础知识和资料：'''{texts}'''
                    你可能还会接收到历史对话中的相关问题和答案：'''{his_record}'''
                    请严格参考基础知识资料以及历史记录（如有），回答用户问题。{add_req}
                    如果用户在找图片，且基础资料中存在该图片，直接返回图片的名字。
                """        
                
    elif task=='filling-tb': # when generating response using KB
        temperature = 0
        top_p = 0.01

        prompt = f"""
                你是一名表格填充人员，你将接收基础资料\n：'''{texts}'''
                请根据已有资料，找到对应查询关键词'''{query}'''的答案。注意以下事项：
                - 如果基础资料中没有包括答案，你的答案是"null"。
                - 注意选择题和打☑的项。
                - 必须返回JSON字典格式，键为"answer"，值为你找到的答案。
                - 答案要直接简单，不要添加任何解释，不要输出任何其他信息。
            """
    
    elif task=='if-history':
        temperature = 0.01
        top_p = 0.01
        
        example_true = '{"trace_history":"是"}'
        example_false = '{"trace_history":"否"}'
        examples_full = '''
                        上一轮提问是：360安全大模型的技术方案包括哪些内容？
                        当前提问是：其中第二项技术方案的具体实现方式是是什么？
                        你应该判断当前提问是上一轮提问的追问，因此回答 {"trace_history":"是"}

                        上一轮提问是：我们单位的报销要求是什么？
                        当前提问是：我可以住多少的酒店？？
                        你应该判断当前提问是上一轮提问的追问，因此回答 {"trace_history":"是"}
                                                
                        上一轮提问是：基坑旁边是否可以存放施工材料？
                        当前提问是：室内管道的最低标高要求是多少？
                        你应该判断当前提问不是上一轮提问的追问，因此回答 {"trace_history":"否"}
        '''
        histories = kwargs['paras']['histories']
        if len(histories)>0:
            last_query = histories[-1][0]
        else:
            last_query = ''
            
        prompt = f"""
                你会收到当前用户的提问'''{query}'''
                你还会收到上一轮的用户问答记录'''{his_record}'''
                请判断，当前提问'''{query}'''是否为在上一轮问答的基础上的追加问题。
                你的回答为JSON字典格式，键为"trace_history"，值只能为"是"或者"否"，例如{example_true}或者{example_false}
                不要输出其他内容，回答不要超过5个字。下面给出了一些例子：
            """
        prompt = prompt + examples_full
        
    elif task=='summary':
        try:
            limit = kwargs['paras']['out_limit']
            add_req = kwargs['paras']['add_req']
        except:
            limit = int(len(texts)*0.2)
            add_req = ''
        prompt = f"""
                你将接收一段文字材料：{texts}
                你需要总结以上文字材料主要内容，字数不超过{limit}字。{add_req}
                你的回答只包含提炼后的文本，不要返回其他多余信息。
            """
    elif task=='summary-keywords':
        max_tokens = kwargs['paras']['out_limit']
        kw_num = kwargs['paras']['kw_number']

        example = '''
            {"keywords":"关键词1;关键词2;关键词3"}
        '''
        
        prompt = f"""
                提取下面这段文字的主题关键词，数量不超过{kw_num}个.
                你的回答为JSON字典格式，键为"keywords"，值为你提取的关键词，关键词要反映文本主题，关键词之间用分号";"隔开，例如 {example}
                '''{texts}'''
            """
    
    elif task=='analyze-paras':
        prompt = f"""
                对下面这段输入的文本内容进行分析：'''{texts}'''
                检查文本内容是否可基于一些明显标识，可将内容分为若干段落，这些明显标识指各类数字类标识，比如“1、”和“2、”等。
                你的回答只能是json形式，键是"answer"，值是"true"或"false"。
                除以上json返回外，不要返回其他多余信息。
            """

    elif task=='rewrite-paras':
        max_tokens = kwargs['paras']['out_limit']
        add_req = kwargs['paras']['add_req']
        model = 'qwen-max'
        
        prompt =  f"""
                你的任务是，润色并重新生成以下段落，不超过{int(1.3*max_tokens)}字，与原文语义和格式相符：
                '''{texts}'''
                
                生成过程需参考上文（如有）]：
                '''{query}'''

                生成内容忽略参考上文中的分点论述，例如 "1) a. 一、" 等
                生成内容不要使用MARKDOWN符号。
                {add_req}
            """

    elif task=='rewrite-paras-whole':
        max_tokens = kwargs['paras']['out_limit']
        add_req = kwargs['paras']['add_req']
        model = 'qwen2-72b-instruct'
        
        prompt =  f"""
                你的任务是，重新生成并润色以下内容，不超过{int(1.3*max_tokens)}字，与原文语义和格式相符，注意使用句间停顿，避免输出无停顿长句：
                '''{texts}'''
                
                生成的内容要承接上文（如有），但不要包括其中内容，避免重复：
                上文：
                '''{query}'''

                生成的内容要使用1. a. 等形式，分层分点论述，也要参考上文的分点论述形式。
                注意数字形式、公式等均与参考上文一致。
                生成内容不要使用MARKDOWN符号。
                {add_req}
            """

    elif task=='abridge-paras':
        topic = kwargs['paras']['topic']
        max_tokens = kwargs['paras']['out_limit']
        model = 'qwen1.5-110b-chat'
        
        prompt =  f"""
                对下面这段输入的文本内容进行提炼：'''{texts}'''
                在提炼过程中，你要做到以下几点：
                1）检查文本中可用于分段的明显标识，比如"1)"、"2)"等数字或者"a)"、"b)"等字母，按这些标识将文本分段分行。
                2）提炼要符合主题：{topic}（如有）。
                3）提炼后的文本体量要比原文缩减，不要有任何解释。
            """
        
    elif task=='extension-paras':
        topic = kwargs['paras']['topic']
        max_tokens = kwargs['paras']['out_limit']
        model = 'qwen1.5-110b-chat'
        
        prompt =  f"""
                对下面这段输入的文本内容进行扩写：'''{texts}'''
                在扩写过程中，你要做到以下几点：
                1）检查文本中可用于分段的明显标识，比如"1)"、"2)"等数字或者"a)"、"b)"等字母，按这些标识将文本分段分行。
                2）扩写要符合主题：{topic}（如有）。
                3）扩写后的文本体量要比原文本增加，不要有任何解释。
            """
        
    elif task=='summary-plain':
        temperature = 0.01
        max_tokens = kwargs['paras']['out_limit']
        
        prompt = f"""
                对下面这段文本总结提炼为一段精简的摘要，保留其关键含义，不要超过{max_tokens}字。
                '''{texts}'''
            """
        
    elif task=='gen-ques':
        temperature = 0.5
        top_p = 0.5
        
        prompt = f"""
                你的角色是一名专为建设工程领域编制专业考试试卷的考官
                你会接收到一段知识内容：'''{texts}'''
                你的任务是基于这段内容出一道问答题，注意：
                题干要能适合用于专业考试（如一级建造师、二级建造师等）；题干的表述要具体详细，且符合给定的知识内容，不要超过50个字
            """
            
    elif task =='rerank':
        temperature = 0
        top_p = 0.01
        model = 'qwen-max'
        
        length = kwargs['paras']['length']
        example = '''
                输入用户提问：'提问'
                输入列表：'[知识路径1, 知识路径2, 知识路径3]'
                你的返回:
                {
                    知识路径2 : 0
                    知识路径1 : 1
                    知识路径3 : 2
                }
            '''
        prompt = f"""
                你会接收到一个用户提问：'''{query}'''
                你还会接收到一个列表，其包含多个可能与这个提问相关的知识文件的存储路径：'''{texts}'''
                
                你的任务是判断列表中每个知识文件与用户提问的相关性和匹配性，并按相关性从高到低，对列表中所有知识文件路径重新排序。
                相关性最高的知识路径，其排序越靠前（最小为0），相关性越低的路径，其排序月靠后（最大为列表长度{length}）
                一个例子如下所示：
                {example}
                你返回的只能是一个JSON格式的字典，该字典的键为输入列表中的知识文件路径，值为其排序。
            """
            
    elif task=='parse-titles': # this one is NOT supported by API currently
        temperature = 0
        top_p = 0.01
        
        example = '''一个输出例子如下：
                {
                  "ROOT": {
                    "1 发展现状": {
                      "1.1 市场情况": {},
                      "1.2 政策情况": {},
                      "1.3 技术情况": {}
                      },
                    
                    "2研究内容": {
                      "2.1 研究内容1": {},
                      "2.2 研究内容2": {}
                    }
                    
                    "技术现状": {}
                  }
                }
            '''
        instruct = f"""
                现在，请根据上述指令和示例，解析下面的输入标题：'''{texts}'''
            """
        prompt = f"""
                你会收到系列文档的标题
                你需要将输入的标题转化为JSON字典的树状结构，父节点是高层级标题，子叶节点是低层级的标题
                你要重点关注标题形式，同级的标题具有相同的形式（比如大小括号、中英文数字的使用等）
                注意标题JSON结构中的节点只能是给定的先验标题，不要自己生成新的标题，也不要对原标题字符进行任何修改或拆分。
                你的返回只能是JSON字典，该JSON结构的根节点固定为ROOT，所有其他节点的键是标题名称，值是下层节点；如下层没有其他标题，则其值固定为空字典。
                一个例子如下所示：
            """
        prompt = prompt + example + instruct 
    
    elif task=='judge-answer':
        other_texts = kwargs['paras']['other_texts']
        prompt = f"""
                你会收到一个问题：{query}
                你还会收到两个答案。第一个是标准答案：{texts}；第二个是待评估的答案：{other_texts}
                你需要判断待评估答案与标准答案的符合程度，给出0-10的整数打分，打分越高，代表待评估答案越符合标准答案。
                注意，待评估答案不需要和标准答案完全一致，但关键数据和表达的意思必须符合。
                你的返回只能是JSON形式，其中键为"score"，值为你的打分。
            """
        keys_ = ['answer']
    
    elif task=='financial_test':
        temperature = 0
        top_p = 0.01
        max_tokens = 500
        model = 'qwen-max'

        prompt= f'''
        你是一名上市公司财务审计专家，现需对目标公司上传的财务报告执行反欺诈核查。

        # 任务
        请基于但不限于以下一项或多项分析维度和财务报表，从财报中逆向推导风险信号，并输出结构化核查结论。你可以在<report></report>标签内获得财务报表。

        # 分析纬度
        1. 数据一致性审查
        横向对比：评估财报内各科目（如收入、成本、费用、资产、负债等）之间的勾稽关系是否合理，是否存在逻辑矛盾。例如，检查收入增长但现金流下降且无合理解释的情况。
        纵向对比：通过对比多期财报中的关键财务指标（如营收、毛利率、现金流、应收账款周转率、存货周转率、大客户/供应商占比等），识别不符合行业周期规律的异常波动或其他“反常识”现象。
        2. 附注信息交叉验证
        会计政策与估计变更分析：检查公司是否频繁调整会计政策或估计方法（如折旧年限、坏账准备计提比例），以掩盖真实财务状况。
        关联交易披露：评估关联方交易披露的透明度，核对重大关联交易披露与业务逻辑的一致性，分析是否存在通过关联方交易转移利润或虚增收入的行为。
        表外负债核查：分析或有负债、担保事项的披露完整性和潜在风险敞口。
        3. 管理层讨论分析
        归因逻辑检验：结合管理层讨论与分析（MD&A）部分，检查其表述是否与财务数据逻辑自洽。例如，管理层宣称“市场扩张”但销售费用率却下降的情况。
        4. 其他财务造假的分析维度
        根据具体财报内容，补充其他可能涉及的分析维度。

        # 财务报表
        <report>
        {texts}
        </report>

        仅分析上述报告，不要输出其他无关信息
    '''

    elif task=='financial_judge':
        temperature = 0
        top_p = 0.01
        max_tokens = 500
        model = 'qwen-max'

        prompt= f'''
        你是一名上市公司财务审计专家，现需对目标公司上传的财务报告执行反欺诈核查。

        # 任务
        请基于财务报表信息，判断{query}所描述的“股市小作文”的可信度，你可以在<report></report>标签内获得财务报表。

        # 财务报表
        <report>
        {texts}
        </report>

        仅分析上述报告，不要输出其他无关信息
    '''
        
    else:
        pass

    limit_len = 20000
    prompt = prompt[:limit_len] + '...' if len(prompt) > limit_len else prompt
    return prompt, model, temperature, top_p, max_tokens, api_key


def call_qwen(model='qwen-turbo', **kwargs):
    def get_answer(prompt, model, temperature, top_p, max_tokens, api_key, **kwargs):
        logger.debug('prompt_len:{}', len(prompt))
        response = call_with_messages(prompt, model, temperature, top_p, max_tokens, api_key) # actual interaction with QWEN server
        try:
            if response and response.status_code == HTTPStatus.OK:
                output = response.output
                if output and output.get("choices"):
                    message = output["choices"][0].get("message")
                    if message:
                        answer = eval_response(message.get("content"), config=kwargs['config'])
                        logger.debug('answer:{}', answer)
            else:
                answer = None
        except:
            answer = None
            print(response)
            logger.debug('not answer, response:{}', response)
        return answer

    prompt, model, temperature, top_p, max_tokens, api_key = build_qwen_prompt(model, **kwargs)
    answer = get_answer(prompt, model, temperature, top_p, max_tokens, api_key, **kwargs)
    return answer


def call_qwen_stream(model='qwen-turbo', **kwargs):
    SYSTEM_ROLE = '你是一位得力的助手，拥有渊博的行业知识。'
    def chat_qwen(prompt, model, temperature, top_p, max_tokens, qwen_apikey, stream=False, histories=[]):
        # logger.debug('prompt_len:{}', len(prompt))
        messages=[{'role': 'system', 'content': SYSTEM_ROLE}, ]
        for history in histories:
            messages.append({"role": "user", "content": history[0]})
            messages.append({'role': 'assistant', 'content': history[1]})
        messages.append({ "role" :"user" , "content" : prompt})
        base_url = 'https://dashscope.aliyuncs.com/compatible-mode/v1'
        client = OpenAI(api_key=qwen_apikey, base_url=base_url)
        response_ = client.chat.completions.create(model=model, 
                                                    messages=messages, 
                                                    max_tokens=max_tokens, # could be annotated?
                                                    top_p=top_p,
                                                    extra_body={'repetition_penalty': 1.2},
                                                    temperature=temperature,
                                                    stream=stream,
                                                    stream_options={"include_usage": True})
        if not stream:
            answer = response_.choices[0].message.content
            return answer
        
        res = ''
        old_ch = ''
        for chunk in response_:
            if chunk.usage:
                logger.info('prompt_tokens:{}, completion_tokens:{}, total_tokens:{}', chunk.usage.prompt_tokens, 
                             chunk.usage.completion_tokens, chunk.usage.total_tokens)
                break
            res_text = chunk.choices[0].delta.content
            if res_text: # \n后面有空格会导致前端处理异常
                new_text = ''
                for ch in res_text:
                    if old_ch == '\n' and ch == ' ':
                        continue
                    new_text += ch
                    old_ch = ch
                res += res_text
                data = json.dumps({'text':new_text}, ensure_ascii=False)
                yield f'data: {data}\n\n'
        logger.debug('res:{}', res)

    prompt, model, temperature, top_p, max_tokens, api_key = build_qwen_prompt(model, **kwargs)
    yield from chat_qwen(prompt, model, temperature, top_p, max_tokens, api_key, stream=True)